﻿#include <algorithm>
#include <gtest/gtest.h>
#include <iostream>

#include <xy/solver/diff.h>
#include <xy/utils/output.h>

using namespace xy;
using namespace xy::math;
using namespace xy::solver;

#define EXP 2.718281828

TEST(TestSolverDiff, ODE)
{
    double lambda = -100;
    double eta = 1.5;
    // double k = 0.01;
    // double T = 1;

    // // 需要注意 N = T / k 的计算方法可能不准
    // int N = 100;

    // 应该这样计算
    double T = 1;

    auto f = [lambda](vecxd x, double t) { return vecxd{lambda * (x[0] - cos(t)) - sin(t)}; };
    auto u = [lambda, eta](double t) { return pow(EXP, lambda * t) * (eta - 1) + cos(t); };

    {
        std::vector<double> err(3);

        int N = 160;
        for (int i = 0; i < err.size(); i++)
        {
            // 注意要多给一个
            auto method = std::make_shared<adams_bashforth>(2, T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method);
            err[i] = fabs(G.back()[0] - u(T));
            N *= 2;
        }
        std::cout << "adams_bashforth" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }

    {
        std::vector<double> err(3);

        int N = 80;
        for (int i = 0; i < err.size(); i++)
        {
            // 注意要多给一个
            auto method = std::make_shared<adams_moulton>(2, T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method);
            err[i] = fabs(G.back()[0] - u(T));
            N *= 2;
        }
        std::cout << "adams_moulton" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }

    {
        std::vector<double> err(3);

        int N = 80;
        for (int i = 0; i < err.size(); i++)
        {
            // 注意要多给一个
            auto method = std::make_shared<bdfs>(2, T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method);
            err[i] = fabs(G.back()[0] - u(T));
            N *= 2;
        }
        std::cout << "bdfs" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }

    {
        std::vector<double> err(3);

        int N = 80;
        for (int i = 0; i < err.size(); i++)
        {
            // 注意要多给一个
            auto method = std::make_shared<classical_RK4>(T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method, false);
            err[i] = fabs(G.back()[0] - u(T));
            EXPECT_LE(fabs(G.back()[0] - u(T)), 1e-5);
            N *= 2;
        }
        std::cout << "classical_RK4" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }

    {
        std::vector<double> err(3);

        int N = 100;
        for (int i = 0; i < err.size(); i++)
        {
            auto method = std::make_shared<ESDIRK6>(T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method, false);
            err[i] = fabs(G.back()[0] - u(T));
            EXPECT_LE(fabs(G.back()[0] - u(T)), 1e-4);
            N *= 2;
        }
        std::cout << "ESDIRK6" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }

    {
        std::vector<double> err(3);

        int N = 10;
        for (int i = 0; i < err.size(); i++)
        {
            auto method = std::make_shared<gauss_legendre_RK>(3, T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method, false);
            err[i] = fabs(G.back()[0] - u(T));
            EXPECT_LE(fabs(G.back()[0] - u(T)), 1e-6);
            N *= 2;
        }
        std::cout << "gauss_legendre_RK" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }

    {
        std::vector<double> err(3);

        int N = 100;
        for (int i = 0; i < err.size(); i++)
        {
            auto method = std::make_shared<fehlberg6>(T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method, false);
            err[i] = fabs(G.back()[0] - u(T));
            EXPECT_LE(fabs(G.back()[0] - u(T)), 1e-6);
            N *= 2;
        }
        std::cout << "fehlberg6" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }

    {
        std::vector<double> err(3);

        int N = 100;
        for (int i = 0; i < err.size(); i++)
        {
            auto method = std::make_shared<dormand_prince7>(T / N);
            time_grid G(N + 1, vecxd{eta});
            ode::solve(f, G, method, false);
            err[i] = fabs(G.back()[0] - u(T));
            EXPECT_LE(fabs(G.back()[0] - u(T)), 1e-5);
            N *= 2;
        }
        std::cout << "dormand_prince7" << std::endl;
        std::cout << "err = [" << err << "];" << std::endl;

        std::vector<double> ratio(2);
        for (int i = 0; i < ratio.size(); i++)
            ratio[i] = err[i] / err[i + 1];
        std::cout << "ratio = [" << ratio << "];" << std::endl;
    }
}

TEST(TestSolverDiff, PDE)
{
    // r = k * nu / h^2
    double nu = 1;
    double r = 1;

    // 空间和时间步长
    double h = 1.0 / 64;
    double k = r * h * h / nu;

    // 空间和时间规模
    int SN = 64;
    int TN = 10;
    double T = TN * k;

    // 边界条件
    auto g = [SN](double t) {
        vecxd v(SN + 1);
        v.set_zero();
        return v;
    };

    // 初值条件
    auto u = [](double x) {
        double r = 0;
        if (9.0 / 20 <= x && x < 0.5)
            r = 20 * (x - 9.0 / 20);
        if (0.5 <= x && x < 11.0 / 20)
            r = -20 * (x - 11.0 / 20);
        return r;
    };

    // 时间网格
    time_grid G(TN + 1, vecxd(SN + 1));

    // 填充初值
    auto &v = G[0];
    for (int i = 0; i < SN + 1; i++)
    {
        double x = i * h;
        v[i] = u(x);
    }

    // 系数矩阵 A
    matxd A(SN + 1, SN + 1);
    A.set_zero();
    int M = SN + 1;
    for (int i = 1; i < M - 1; i++)
    {
        A(i, i) = -2 * nu / h / h;
        A(i, i + 1) = nu / h / h;
        A(i, i - 1) = nu / h / h;
    }
    A(0, 0) = A(M - 1, M - 1) = -2 * nu / h / h;
    A(0, 1) = A(M - 1, M - 2) = nu / h / h;

    auto method = std::make_shared<ESDIRK6>(k);
    pde::solve(A, g, G, method);

    std::cout << "U = [" << G.back() << "];" << std::endl;
}